import numpy as np
import torch
import model
import utils

# 实例化模型
context_encoder = model.ContextualEncoder()
decoder = model.WaveFieldDecoder()

# 加载保存的参数
context_encoder.load_state_dict(torch.load("./net/OW-DGM_context_encoder_model.pth"))
decoder.load_state_dict(torch.load("./net/OW-DGM_decoder_model.pth"))

# 设置为评估模式
context_encoder.eval()
decoder.eval()
all_generated_wave_fields = []

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("loading buoy data........")
buoy_data = np.load('data/buoy_data/buoy_data_test.npy')
buoy_data = torch.tensor(buoy_data)

# 初始化推理模块
inference = utils.Inference(context_encoder, decoder, device)

# 创建数据集和数据加载器
dataset = utils.BuoySeriesDataset(buoy_data, batch_size=128)
dataloader = utils.DataLoader(dataset, batch_size=None, shuffle=False)
for batch_idx, buoy_data in enumerate(dataloader):
    print(f"Processing batch {batch_idx + 1}/{len(dataloader)}...")
    generated_wave_fields = inference.generate_wave_field(buoy_data)
    all_generated_wave_fields.append(generated_wave_fields)

# 将所有生成的波场拼接为一个张量
all_generated_wave_fields = torch.cat(all_generated_wave_fields, dim=0)
print("所有生成波场的形状:", all_generated_wave_fields.shape)
generated_wave_fields = all_generated_wave_fields.cpu().numpy()
np.save('data/generated_wave/generated_wave_fields.npy', generated_wave_fields)
print("数据已成功保存为.npy文件")
